from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QFileDialog
import math
import numpy as np
import cv2
class Extract(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()
        self.setObjectName("Extract")
        self.resize(1094, 508)
        self.gridLayout_2 = QtWidgets.QGridLayout(self)
        self.gridLayout_2.setContentsMargins(11, 11, 11, 11)
        self.gridLayout_2.setSpacing(6)
        self.gridLayout_2.setObjectName("gridLayout_2")
        self.horizontalLayout = QtWidgets.QHBoxLayout()
        self.horizontalLayout.setSpacing(6)
        self.horizontalLayout.setObjectName("horizontalLayout")
        self.play_button = QtWidgets.QPushButton(self)
        self.play_button.setMinimumSize(QtCore.QSize(131, 41))
        self.play_button.setObjectName("play_button")
        self.horizontalLayout.addWidget(self.play_button)
        self.start_button = QtWidgets.QPushButton(self)
        self.start_button.setMinimumSize(QtCore.QSize(131, 41))
        self.start_button.setObjectName("start_button")
        self.horizontalLayout.addWidget(self.start_button)
        self.lineEdit = QtWidgets.QLineEdit(self)
        self.lineEdit.setMinimumSize(QtCore.QSize(600, 0))
        self.lineEdit.setMaximumSize(QtCore.QSize(1000, 16777215))
        self.lineEdit.setReadOnly(True)
        self.lineEdit.setObjectName("lineEdit")
        self.horizontalLayout.addWidget(self.lineEdit)
        self.gridLayout_2.addLayout(self.horizontalLayout, 0, 0, 1, 1)
        self.graphicsView = QtWidgets.QGraphicsView(self)
        self.graphicsView.setObjectName("graphicsView")
        self.gridLayout_2.addWidget(self.graphicsView, 1, 0, 1, 1)

        self.data_file = None
        self.width = 384
        self.height = 288
        self.frame_size = self.width * self.height * 2  # 2 bytes per pixel
        self.frames = 0
        self.image = []
        self.retranslateUi()

    def retranslateUi(self):
        self.setWindowTitle("Extract")
        self.play_button.setText("Start Extracting")
        self.play_button.clicked.connect(self.toggle_playback)
        self.start_button.setText("Choose file.data")
        self.start_button.clicked.connect(self.choose_data_file)
        self.scene = QtWidgets.QGraphicsScene()

    def choose_data_file(self):
        options = QFileDialog.Options()
        options |= QFileDialog.DontUseNativeDialog
        file_path, _ = QFileDialog.getOpenFileName(self, "Open .data File", "", "Data Files (*.data);;All Files (*)",
                                                   options=options)
        if file_path:
            self.lineEdit.setText(file_path)
            self.data_file = open(file_path, 'rb')
            self.data_file.seek(0, 2)  # Move to end of file
            file_size = self.data_file.tell()
            self.data_file.seek(0)  # Move back to start
            self.frames = file_size // self.frame_size

    def toggle_playback(self):
        if self.data_file:
            self.image.clear()
            self.update_image()
        else:
            QtWidgets.QMessageBox.warning(self, "No File", "Please choose a .data file first.")

    def get_image(self):
        for i in range(self.frames):
            # print(i)
            self.data_file.seek(i * self.frame_size)
            frame_data = self.data_file.read(self.frame_size)
            image = np.frombuffer(frame_data, dtype=np.uint16).reshape(self.height, self.width)
            if image.var() > 30000:
                # print(image.var())
                self.image.append(image)
        if self.data_file:
            self.data_file.close()

    def combine_images(self):
        cols = 5
        rows = math.ceil(len(self.image) / cols)
        combined_image = np.zeros((288 * rows, 384 * cols), dtype=np.uint8)
        if len(self.image) > 0:
            h, w = self.image[0].shape
            for i, img in enumerate(self.image):
                image = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
                r = i // cols
                c = i % cols
                combined_image[r * h:(r + 1) * h, c * w:(c + 1) * w] = image
        return combined_image

    def update_image(self):
        self.get_image()
        combined_image = self.combine_images()
        h, w = combined_image.shape
        q_image = QtGui.QImage(combined_image.data, w, h, w, QtGui.QImage.Format_Grayscale8)
        item = QtWidgets.QGraphicsPixmapItem(QtGui.QPixmap.fromImage(q_image))
        # 清空场景并添加新的 QPixmapItem
        self.scene.clear()
        self.scene.addItem(item)

        self.graphicsView.setScene(self.scene)
        self.graphicsView.setDragMode(QtWidgets.QGraphicsView.ScrollHandDrag)
        # self.graphicsView.setTransformationAnchor(QtWidgets.QGraphicsView.AnchorUnderMouse)
        self.data_file.close()
        self.data_file = None

    def closeEvent(self, event):
        if self.data_file and not self.data_file.closed:
            self.data_file.close()
        event.accept()

if __name__ == "__main__":
    import sys
    app = QtWidgets.QApplication(sys.argv)
    Widget = Extract()
    Widget.show()
    sys.exit(app.exec_())
